package me.chengshu.springai.chatclient.controller;

import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import me.chengshu.springai.chatclient.advisors.ReReadingAdvisor;
import me.chengshu.springai.chatclient.entity.DynamicModelAndOptions;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SafeGuardAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.deepseek.DeepSeekChatModel;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.List;

@RequestMapping("/chat")
@RestController
public class ChatClientController {

    HashMap<String, ChatModel> chatModels = new HashMap<>();

    public ChatClientController(DashScopeChatModel dashScopeChatModel,
                                DeepSeekChatModel deepSeekChatModel,
                                OllamaChatModel ollamaChatModel) {
        // 将 ChatModel 对象添加到 chatModels 中，并使用模型名称作为键
        chatModels.put("deepseek", deepSeekChatModel);
        chatModels.put("ollama", ollamaChatModel);
        chatModels.put("dashscope", dashScopeChatModel);
    }

    /**
     * 在多平台中选择指定模型
     * @param content 用户输入内容
     * @param platform 平台
     * @param model 模型名称
     * @param temperature 模型温度
     * @return 响应流
     */
    @RequestMapping(value = "/one/in/more", produces = "text/stream;charset=UTF-8")
    public Flux<String> chatOneInMoreModel(@RequestParam String content,
                                           @RequestParam String platform,
                                           @RequestParam Double temperature,
                                           @RequestParam String model){
        ChatModel chatModel = chatModels.get(platform);
        // TODO: 校验参数
        ChatOptions options = ChatOptions.builder()
                .temperature(temperature)
                .model(model)
                .build();
        ChatClient chatClient = ChatClient.builder(chatModel)
                // .defaultSystem("") 系统提示词: 你是什么, 你能做什么? 你要注意什么? 等等
                .defaultAdvisors(
                        new SimpleLoggerAdvisor(),
                        new ReReadingAdvisor(),
                        new SafeGuardAdvisor(List.of("CEO"))
                ) // 日志拦截器
                .defaultOptions(options)
                .build();
        return chatClient.prompt()
                .system(promptSystemSpec -> promptSystemSpec.param("name", "me"))
                .user(content)
                .stream()
                .content();
    }
}
